from ..core.problem_definition_base import ProblemDefinitionBase
from ..core.simulation_parameters import SimulationParameters
from typing import Dict, Any, Optional, List

# FEniCSx相关的导入，同样使用条件导入
try:
    import dolfinx
    from dolfinx import fem, mesh
    import ufl
    from ufl.core.expr import Expr as UFLExpr # 显式导入并重命名以避免冲突
    from mpi4py import MPI
    FENICSX_AVAILABLE = True
except ImportError:
    FENICSX_AVAILABLE = False
    # 用于类型提示的虚拟类型 
    class fem:
        class Function: pass
        class DirichletBC: pass
    class mesh:
        class MeshTags: pass
    class ufl:
        class Form: pass
        class Argument: pass
        class Measure: pass
    class UFLExpr: pass # 虚拟类型 UFLExpr
    class MPI:
        class Comm: pass

class WaterTransportProblem(ProblemDefinitionBase):
    """
    求解水传输的问题定义，通常在离聚物相中
    (例如，水含量 'lambda' 或水活度 'a_w')。
    这涉及扩散和电渗拖曳 (EOD)。
    """
    def __init__(self,
                 problem_name: str, # 例如, "lambda_pem", "aw_cathode_cl"
                 fem_entities: Dict[str, Any],
                 material_functions: Dict[str, Any],
                 params: SimulationParameters,
                 # 可选的 UFL 表达式，用于从其他物理场耦合的项
                 proton_current_density_vector_ufl: Optional[UFLExpr] = None, # 质子电流密度矢量 j_H+ (A/m^2)
                 volumetric_water_source_term_ufl: Optional[UFLExpr] = None, # 体积水合物源项 (mol/m^3/s 或 kg/m^3/s)
                 domain_tags: Optional[mesh.MeshTags] = None,
                 dx_measure: Optional[ufl.Measure] = None,
                 ds_measure: Optional[ufl.Measure] = None,
                 boundary_markers: Optional[Dict[str, Any]] = None,
                 comm: Optional[MPI.Comm] = None,
                 **kwargs):
        """
        初始化水传输问题。

        参数:
            problem_name (str): 此水传输问题的唯一名称。
            fem_entities (Dict[str, Any]): 基本的 FEniCSx 实体。
            material_functions (Dict[str, Any]): 材料属性，期望包含有效水
                                                扩散系数 ('water_diffusivity_ionomer') 和
                                                电渗拖曳系数 ('electro_osmotic_drag_coeff')。
            params (SimulationParameters): 所有模拟参数。
            proton_current_density_vector_ufl (Optional[UFLExpr]): 
                质子电流密度矢量 (j_H+) 的 UFL 表达式，EOD 项需要。
            volumetric_water_source_term_ufl (Optional[UFLExpr]):
                体积水合物源/汇项的 UFL 表达式 (例如，来自电化学反应
                或相变，如果在此处建模为源/汇项)。
            domain_tags (Optional[dolfinx.mesh.MeshTags]): 域的网格标签。
            dx_measure (Optional[ufl.Measure]): 用于体积积分的 UFL 度量。
            ds_measure (Optional[ufl.Measure]): 用于面积积分的 UFL 度量。
            boundary_markers (Optional[Dict[str, Any]]): 将边界名称映射到标记 ID。
            comm (Optional[MPI.Comm]): MPI 通信器。
            **kwargs: 额外的关键字参数。
        """
        self.proton_current_density_vector_ufl = proton_current_density_vector_ufl
        self.volumetric_water_source_term_ufl = volumetric_water_source_term_ufl
        # fem_entities 中的变量键，例如 "lambda_pem" 或 problem_name
        self.variable_key_in_fem_entities = kwargs.pop('variable_key', problem_name)

        super().__init__(problem_name=problem_name,
                         fem_entities=fem_entities,
                         material_functions=material_functions,
                         params=params,
                         domain_tags=domain_tags,
                         dx_measure=dx_measure,
                         ds_measure=ds_measure,
                         boundary_markers=boundary_markers,
                         comm=comm,
                         **kwargs)

    def setup_problem(self) -> None:
        """
        为水传输 (例如，水含量 lambda) 设置变分问题。

        离聚物中水含量的常见控制方程 (lambda):
            ε_m * (ρ_m_dry / EW) * ∂λ/∂t + ∇ ⋅ N_wat = S_wat
        对于稳态 (∂λ/∂t = 0):
            ∇ ⋅ N_wat = S_wat
        
        水通量 N_wat (mol/m^2/s) 可以描述为 (Springer 等人 / Weber 等人):
            N_wat = N_drag + N_diff
                  = (n_d / F) ⋅ j_H+  - D_wat_eff ⋅ ∇C_wat
        如果求解 lambda, C_wat = (ρ_m_dry / EW) * lambda, 因此 ∇C_wat = (ρ_m_dry / EW) * ∇lambda。
            N_wat = (n_d / F) ⋅ j_H+  - D_lambda_eff ⋅ ∇lambda
        其中 D_lambda_eff = (ρ_m_dry / EW) * D_wat_eff。

        因此，稳态方程变为:
            ∇ ⋅ [ (n_d / F) ⋅ j_H+  - D_lambda_eff ⋅ ∇λ ] = S_wat
            -∇ ⋅ (D_lambda_eff ⋅ ∇λ) + ∇ ⋅ ((n_d / F) ⋅ j_H+) = S_wat

        变分形式 (假设 D_lambda_eff 不是 lambda 的函数，对于线性形式):
            ∫_Ω (D_lambda_eff ⋅ ∇trial_λ ⋅ ∇test_λ) dΩ - ∫_Ω ( (n_d/F) ⋅ j_H+ ⋅ ∇test_λ ) dΩ = ∫_Ω S_wat ⋅ test_λ dΩ + BC_terms
        EOD 项 (∇ ⋅ ((n_d/F) ⋅ j_H+)) 中的 j_H+ 如果已知，可以通过分部积分写为 - (n_d/F) ⋅ j_H+ ⋅ ∇test_λ 。
        (原先的注释说 + ∫ ... (n_d/F) j_H+ ⋅ ∇test_λ dΩ，这里根据 N_wat 的定义和散度定理，应该是负号，
         因为它是从 N_wat 散度项的第二部分 ∇⋅( (n_d/F) j_H+ λ ) 来的，如果使用标准方法处理对流项，则符号取决于如何移动到另一侧，
         或者直接从 ∇⋅N_drag = ∇⋅( (n_d/F) j_H+ ) 构建，并用试函数相乘积分。) 
        更准确地说，如果 N_drag = (n_d/F) * j_H+，那么它的散度贡献是 ∇⋅N_drag。
        ∫ (∇⋅N_drag) test_λ dΩ = - ∫ N_drag ⋅ ∇test_λ dΩ + ∫_Γ N_drag⋅n test_λ dS.
        所以，在体积积分中是 -∫ (n_d/F) j_H+ ⋅ ∇test_λ dΩ。这会出现在线性形式中，如果 j_H+ 是已知的。
        如果 n_d 也依赖于 lambda (试函数)，则该项将贡献给双线性形式。
        这里假设 j_H+ (proton_current_density_vector_ufl) 是已知的。
        EOD 项 -∫ (n_d/F) j_H+ ⋅ ∇test_λ dx 像一个源项，如果 j_H+ 和 n_d 已知。
        或者，如果它写成 ∇⋅(coeff * trial_lambda)，它贡献给双线性形式的对流部分。
        基于 N_wat = (n_d / F) ⋅ j_H+  - D_lambda_eff ⋅ ∇lambda, 
        则 ∇ ⋅ N_wat = ∇⋅((n_d / F) ⋅ j_H+) - ∇⋅(D_lambda_eff ⋅ ∇lambda) = S_wat
        ∫ (- ∇⋅(D_lambda_eff ⋅ ∇lambda) + ∇⋅((n_d / F) ⋅ j_H+) ) test_λ dx = ∫ S_wat test_λ dx
        ∫ (D_lambda_eff ⋅ ∇lambda ⋅ ∇test_λ) dx - ∫_Γ (D_lambda_eff ∇lambda ⋅ n) test_λ dS 
        - ∫ ((n_d / F) ⋅ j_H+) ⋅ ∇test_λ dx + ∫_Γ ((n_d/F)j_H+⋅n) test_λ dS = ∫ S_wat test_λ dx
        双线性形式包含: ∫ (D_lambda_eff ⋅ ∇trial_λ ⋅ ∇test_λ) dx
        线性形式包含 EOD:  ∫ ((n_d / F) ⋅ j_H+) ⋅ ∇test_λ dx + ∫ S_wat ⋅ test_λ dx
        这里，EOD 贡献 ∫ ( (n_d / F) j_H+ ⋅ ∇test_λ) dx 移到左边会变成负号，但它本身是线性形式的一部分。
        如果我们将 EOD 视为对流项，即 N_drag = v_drag * C_water，其中 v_drag ∝ j_H+，那么形如 div(v_drag * trial_lambda)。
        将其移至左侧，并用试函数相乘积分，将得到：∫ (v_drag ⋅ ∇trial_λ) test_λ dx。
        """
        if not FENICSX_AVAILABLE:
            print(f"警告: FEniCSx 不可用。跳过 {self.problem_name} 的 setup_problem()。")
            self._bilinear_form = None # 确保这些在 FEniCSx 不可用时被定义为 None
            self._linear_form = None
            self._bcs = []
            return

        # 1. 识别要求解的 fem.Function (例如, lambda) 及其函数空间
        solution_vars = self.fem_entities.get('solution_variables', {})
        if self.variable_key_in_fem_entities in solution_vars:
            self._variable_to_solve = solution_vars[self.variable_key_in_fem_entities]
            if not isinstance(self._variable_to_solve, fem.Function):
                raise TypeError(f"'{self.variable_key_in_fem_entities}' 必须是 dolfinx.fem.Function 类型。")
            V_space = self._variable_to_solve.function_space
            if V_space is None: raise ValueError(f"'{self.variable_key_in_fem_entities}' 的函数空间为 None。")
        else:
            # 尝试从函数空间字典中获取，如果变量尚未创建
            function_spaces = self.fem_entities.get('function_spaces', {})
            if self.variable_key_in_fem_entities in function_spaces:
                V_space = function_spaces[self.variable_key_in_fem_entities]
                raise KeyError(f"'{self.variable_key_in_fem_entities}' 在 solution_variables 中未找到，虽然在 function_spaces 中找到了。请确保 solution_variable 已创建。")
            else:
                raise KeyError(f"求解变量 '{self.variable_key_in_fem_entities}' (用于水传输) 或其函数空间未找到。"
                               f"可用的 solution_variables: {list(solution_vars.keys())}。 "
                               f"可用的 function_spaces: {list(function_spaces.keys())}。")

        # 2. 识别试函数和试验函数
        self._test_function = ufl.TestFunction(V_space)
        self._trial_function = ufl.TrialFunction(V_space)

        trial_var = self._trial_function # lambda 或 a_w 的试函数
        test_var = self._test_function  # 检验函数

        # 3. 获取材料属性
        # 有效水扩散系数 (可能是 D_lambda_eff 或 D_aw_eff)
        D_wat_eff_key = self.material_functions.get("water_diffusivity_ionomer_key", "water_diffusivity_ionomer")
        D_wat_eff = self.material_functions.get(D_wat_eff_key)
        if D_wat_eff is None:
            # 尝试一个更通用的备用键
            D_wat_eff_generic_key = "water_diffusivity"
            D_wat_eff = self.material_functions.get(D_wat_eff_generic_key)
            if D_wat_eff is None:
                raise ValueError(f"材料属性 '{D_wat_eff_key}' 或 '{D_wat_eff_generic_key}' 未找到。")

        if not hasattr(D_wat_eff, '__mul__'): # 检查是否与 UFL 兼容
             raise TypeError(f"材料 '{D_wat_eff_key if self.material_functions.get(D_wat_eff_key) else D_wat_eff_generic_key}' 与 UFL 不兼容。")

        # 电渗拖曳系数 (n_d)
        n_d_key = self.material_functions.get("electro_osmotic_drag_coeff_key", "electro_osmotic_drag_coeff")
        n_d = self.material_functions.get(n_d_key)
        # 如果忽略 EOD 或以其他方式处理，n_d 可以为 None
        if n_d is not None and not hasattr(n_d, '__mul__'): 
            raise TypeError(f"材料 '{n_d_key}' 与 UFL 不兼容。")

        # 法拉第常数 F，从参数中获取
        if not hasattr(self.params, 'physical_constants') or \
           not hasattr(self.params.physical_constants, 'faraday_constant_C_per_mol'): # 检查正确的字段名
            raise ValueError("法拉第常数 'faraday_constant_C_per_mol' 未在 params.physical_constants 中找到。")
        F_const = self.params.physical_constants.faraday_constant_C_per_mol # 使用正确的字段名

        # 4. 定义变分形式分量
        dx = self.dx_measure
        if dx is None: raise ValueError(f"'{self.problem_name}' 未提供 dx_measure。")
        # ds = self.ds_measure # 如果边界条件需要

        # 扩散项: ∫ D_wat_eff ⋅ ∇trial_var ⋅ ∇test_var dΩ
        self._bilinear_form = D_wat_eff * ufl.dot(ufl.grad(trial_var), ufl.grad(test_var)) * dx
        
        # 初始化线性形式: ∫ 0 ⋅ test_var dΩ
        self._linear_form = ufl.as_ufl(0) * test_var * dx 

        # 电渗拖曳 (EOD) 项:
        # Flux: N_drag = (n_d / F) ⋅ j_H+
        # Term in PDE: - ∇ ⋅ N_drag  (if N_drag is on the LHS of ∇⋅N = S)
        # Weak form: ∫ (-∇ ⋅ N_drag) test_var dx = ∫ N_drag ⋅ ∇test_var dx - ∫_Γ (N_drag ⋅ n) test_var dS
        # So, the volume integral is ∫ (n_d / F) ⋅ j_H+ ⋅ ∇test_var dx, which goes into the linear form (RHS of a = L)
        # Or, if we formulate it as a convection term for trial_var:
        # ∇ ⋅ ( (n_d/F) j_H+ trial_var ) -> leads to a term like ∫ ( (n_d/F) j_H+ ⋅ ∇trial_var ) test_var dx in bilinear form
        # Let's use the convection-like term in the bilinear form for now, assuming j_H+ is a known coefficient vector.
        if self.proton_current_density_vector_ufl is not None and n_d is not None:
            if not hasattr(self.proton_current_density_vector_ufl, 'ufl_shape') or self.proton_current_density_vector_ufl.ufl_shape == (): 
                raise TypeError(f"{self.problem_name} 的 proton_current_density_vector_ufl 不是有效的 UFL 矢量类型。")
            
            # 对流速度等效项 v_eod = (n_d / F_const) * self.proton_current_density_vector_ufl
            eod_convection_coeff = (n_d / F_const) * self.proton_current_density_vector_ufl
            # 添加到双线性形式: ∫ (v_eod ⋅ ∇trial_var) ⋅ test_var dΩ
            self._bilinear_form += ufl.dot(eod_convection_coeff, ufl.grad(trial_var)) * test_var * dx
            # print(f"信息: 已为 {self.problem_name} 添加 EOD (对流) 项。")
        elif n_d is not None and self.proton_current_density_vector_ufl is None:
            print(f"警告: 为 {self.problem_name} 提供了 EOD 系数 'n_d'，但 proton_current_density_vector_ufl 为 None。EOD 项将为零。")

        # 体积水合物源/汇项: ∫ S_wat ⋅ test_var dΩ
        if self.volumetric_water_source_term_ufl is not None:
            if not hasattr(self.volumetric_water_source_term_ufl, '__mul__'):
                raise TypeError(f"{self.problem_name} 的 volumetric_water_source_term_ufl 与 UFL 不兼容。")
            self._linear_form += self.volumetric_water_source_term_ufl * test_var * dx

        # 5. 定义狄利克雷边界条件
        self._bcs = [] # 初始化为空列表
        # 示例: 在气体通道界面处固定水含量/活度
        # gc_interface_marker_name = "cathode_cl_gc_interface" # 假设的边界标记名称
        # water_val_param_name = "water_activity_cathode_gc"   # 假设的参数名称

        # if self.boundary_markers and gc_interface_marker_name in self.boundary_markers and \
        #    hasattr(self.params, 'boundary_conditions') and hasattr(self.params.boundary_conditions, water_val_param_name):

        #     marker_id = self.boundary_markers[gc_interface_marker_name]
        #     fixed_value = getattr(self.params.boundary_conditions, water_val_param_name)
            
        #     # 如果求解的是 lambda，而参数是 a_w，则需要转换，反之亦然。
        #     # 这通常在参数模型或材料属性计算中处理，这里假设 fixed_value 与求解变量兼容。
            
        #     # facet_dim = V_space.mesh.topology.dim - 1
        #     # boundary_facets = dolfinx.mesh.locate_entities_boundary(V_space.mesh, facet_dim, lambda x: x == marker_id)
        #     # fixed_dofs = fem.locate_dofs_topological(V_space, facet_dim, boundary_facets)
        #     # self._bcs.append(fem.dirichletbc(fem.Constant(V_space.mesh, fixed_value), fixed_dofs, V_space))
        # else:
        #     # print(f"调试: 未能为 {self.problem_name} 设置水传输边界条件。检查标记和参数。")
        #     pass
            
        # 检查必要的实体是否已设置
        if self._bilinear_form is None or self._linear_form is None:
            raise RuntimeError(f"'{self.problem_name}' 的变分形式未正确设置。")
            
        # print(f"信息: 问题 '{self.problem_name}' (WaterTransportProblem) 设置完成。")
        # print(f"  双线性形式已组装: {self._bilinear_form is not None}")
        # print(f"  线性形式已组装: {self._linear_form is not None}")
        # print(f"  定义的边界条件数量: {len(self._bcs)}") 